#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    : momentum_month.py
@Time    : 2024/12/01
@Author  : Liuli
@Desc    : 月度技术分析指标计算
"""

import numpy as np
import pandas as pd


def calculate_ma_crossover_signals(price_data, short_period=3, long_periods=[6, 9, 12, 15, 18, 21]):
    """
    计算均线死叉信号
    当3个月均线连续死叉多个长期均线时，提示趋势破位
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        日度资产价格数据，会被转换为月度数据
    short_period : int
        短期均线周期（月），默认3个月
    long_periods : list
        长期均线周期列表（月）
        
    Returns:
    --------
    pd.DataFrame
        死叉信号，1表示死叉，0表示非死叉
    """
    # 将日度数据转换为月度数据
    price_data = price_data.resample('M').last()
    
    # 计算短期均线
    short_ma = monthly_data.rolling(window=short_period).mean()
    
    # 计算各个长期均线的死叉信号
    signals = pd.DataFrame(index=monthly_data.index, columns=monthly_data.columns)
    
    for col in monthly_data.columns:
        # 计算所有长期均线
        long_ma_signals = []
        for period in long_periods:
            long_ma = monthly_data[col].rolling(window=period).mean()
            # 当短期均线位于长期均线下方时，记为死叉
            death_cross = (short_ma[col] < long_ma) & (short_ma[col].shift(1) >= long_ma.shift(1))
            long_ma_signals.append(death_cross)
        
        # 只要有一个长期均线发生死叉，就记为信号
        signals[col] = pd.concat(long_ma_signals, axis=1).any(axis=1).astype(int)
    
    return signals


def calculate_macd_signals(price_data, fast=12, slow=26, signal=9):
    """
    计算MACD信号
    使用月度数据，当MACD处于绿柱且绿柱扩大时，提示趋势破位
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        日度资产价格数据，会被转换为月度数据
    fast : int
        快线周期（月）
    slow : int
        慢线周期（月）
    signal : int
        信号线周期（月）
        
    Returns:
    --------
    pd.DataFrame
        MACD信号，1表示趋势破位，0表示非破位
    """
    # 将日度数据转换为月度数据
    price_data = price_data.resample('M').last()

    signals = pd.DataFrame(index=price_data.index, columns=price_data.columns)
    
    for col in price_data.columns:
        # 计算EMA
        exp1 = price_data[col].ewm(span=fast, adjust=False).mean()
        exp2 = price_data[col].ewm(span=slow, adjust=False).mean()
        
        # 计算MACD线和信号线
        macd = exp1 - exp2
        signal_line = macd.ewm(span=signal, adjust=False).mean()
        
        # 计算MACD柱
        histogram = macd - signal_line
        
        # 当MACD柱为负且比前一月更负时，视为趋势破位信号
        signals[col] = ((histogram < 0) & (histogram < histogram.shift(1))).astype(int)
    
    return signals


def calculate_bollinger_signals(price_data, window=20, num_std=2):
    """
    计算布林带信号
    使用月度数据，当价格从上到下穿越中轨时提示趋势破位，突破下轨时不产生信号
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        日度资产价格数据，会被转换为月度数据
    window : int
        移动平均窗口期（月）
    num_std : int
        标准差倍数
        
    Returns:
    --------
    pd.DataFrame
        布林带信号，1表示趋势破位，0表示非破位
    """
    # 将日度数据转换为月度数据
    price_data = price_data.resample('M').last()

    signals = pd.DataFrame(index=price_data.index, columns=price_data.columns)
    
    for col in price_data.columns:
        # 计算移动平均线（中轨）
        middle_band = price_data[col].rolling(window=window).mean()
        
        # 计算标准差
        std = price_data[col].rolling(window=window).std()
        
        # 计算上下轨
        upper_band = middle_band + num_std * std
        lower_band = middle_band - num_std * std
        
        # 计算信号：
        # 1. 从上到下穿越中轨时产生信号
        cross_middle = (price_data[col] < middle_band) & (price_data[col].shift(1) >= middle_band)
        # 2. 但是如果突破下轨，则不产生信号
        below_lower = price_data[col] < lower_band
        
        signals[col] = (cross_middle & ~below_lower).astype(int)
    
    return signals


def calculate_technical_signals(price_data):
    """
    综合计算技术指标信号
    只要有一个指标提示趋势破位，就认为该资产存在技术面风险，风险预算比例除以5
    
    Parameters:
    -----------
    price_data : pd.DataFrame
        月度资产价格数据
        
    Returns:
    --------
    pd.DataFrame
        综合技术信号和风险调整系数
    """
    # 计算各个技术指标的信号
    ma_signals = calculate_ma_crossover_signals(price_data)
    macd_signals = calculate_macd_signals(price_data)
    bollinger_signals = calculate_bollinger_signals(price_data)
    
    # 综合三个信号：任意一个指标产生信号就视为趋势破位
    technical_signals = (ma_signals | macd_signals | bollinger_signals).astype(int)
    
    # 添加风险调整系数
    risk_adjustment = pd.DataFrame(1.0, index=price_data.index, columns=price_data.columns)
    risk_adjustment[technical_signals == 1] = 0.2  # 当有趋势破位信号时，风险预算除以5
    
    return pd.concat({
        'signals': technical_signals,
        'risk_adjustment': risk_adjustment
    }, axis=1, keys=['signals', 'risk_adjustment'])


def calculate_momentum_signals(df: pd.DataFrame, 
                             lookback_periods: int = 12,
                             vol_lookback: int = 36) -> pd.DataFrame:
    """
    计算时序动量信号
    
    参数:
        df: 包含月度收盘价的DataFrame
        lookback_periods: 动量回看期数(默认12个月)
        vol_lookback: 波动率计算的回看期数(默认36个月)
    
    返回:
        包含动量信号的DataFrame
    """
    # 计算月度收益率
    returns = df.pct_change()
    
    # 计算过去N个月的累积收益
    momentum_returns = (1 + returns).rolling(window=lookback_periods).apply(
        lambda x: np.prod(x) - 1
    )
    
    # 计算历史波动率(使用过去3年月度数据)
    rolling_vol = returns.rolling(window=vol_lookback).std() * np.sqrt(12)
    
    # 计算风险调整后的动量得分
    momentum_score = momentum_returns / rolling_vol
    
    # 生成动量信号
    momentum_signals = pd.DataFrame(index=df.index)
    
    # 趋势破位信号(当动量得分从正转负或从负转正时)
    momentum_signals['trend_reversal'] = (
        (momentum_score.shift(1) > 0) & (momentum_score < 0) |  # 多头转空头
        (momentum_score.shift(1) < 0) & (momentum_score > 0)    # 空头转多头
    )
    
    # 添加动量方向
    momentum_signals['momentum_direction'] = np.sign(momentum_score)
    
    # 添加原始动量得分
    momentum_signals['momentum_score'] = momentum_score
    
    return momentum_signals


def generate_positions(signals: pd.DataFrame, 
                      position_size: float = 1.0) -> pd.DataFrame:
    """
    根据动量信号生成仓位
    
    参数:
        signals: 动量信号DataFrame
        position_size: 单个头寸大小
        
    返回:
        包含目标仓位的DataFrame
    """
    positions = pd.DataFrame(index=signals.index)
    positions['target_position'] = signals['momentum_direction'] * position_size
    
    # 在趋势破位时调整仓位
    positions.loc[signals['trend_reversal'], 'target_position'] = 0
    
    return positions


if __name__ == "__main__":
    # 测试代码
    from WindPy import w
    w.start()
    
    # 获取测试数据
    start = '2020-01-01'
    end = '2024-01-19'
    index = ['H00300.CSI', 'H00905.CSI', 'H00852.SH', 'CBA06501.CS', 'CBA04601.CS',
             'CBA14101.CS', 'CBA20901.CS', 'SP500TR.SPI', 'AU.SHF', 'CU.SHF', 'SC.INE']
    
    # 获取日度数据
    data = w.wsd(index, "close", start, end, "")
    price_data = pd.DataFrame(np.array(data.Data).T, index=data.Times, columns=data.Codes)
    
    # 计算技术信号
    tech_analysis = calculate_technical_signals(price_data)
    
    # 输出结果
    print("\n技术信号汇总：")
    print(tech_analysis['signals'].tail())
    
    print("\n风险调整系数：")
    print(tech_analysis['risk_adjustment'].tail())
    
    # 分析各个指标的信号触发频率
    ma_signals = calculate_ma_crossover_signals(price_data)
    macd_signals = calculate_macd_signals(price_data)
    bollinger_signals = calculate_bollinger_signals(price_data)
    
    print("\n各指标信号触发频率：")
    print("均线系统：", ma_signals.mean().mean())
    print("MACD系统：", macd_signals.mean().mean())
    print("布林带系统：", bollinger_signals.mean().mean()) 